from qiskit import QuantumCircuit, transpile
from qiskit.providers import Backend
import numpy as np
from typing import List, Dict, Union, Optional
import itertools
from qiskit.circuit import QuantumCircuit
from qiskit_aer import AerSimulator
from qiskit.visualization import plot_histogram
from qiskit.transpiler.preset_passmanagers import generate_preset_pass_manager

class BenchmarkCircuit: 
    """
    A class for creating benchmark circuits to characterize readout errors.
    
    This class prepares circuits that initialize qubits in different computational basis states,
    measures them, and uses the results to construct a transition matrix for readout error mitigation.
    """
    def __init__(self, num_qubits: int, qubit_indices: Optional[List[int]] = None):
        """
        Initialize the benchmark circuit generator.
        
        Args:
            num_qubits: Total number of qubits in the circuit
            qubit_indices: Specific qubits to characterize. If None, all qubits are used.
        """
        self.num_qubits = num_qubits
        
        # If qubit_indices is not provided, use all qubits
        if qubit_indices is None:
            self.qubit_indices = list(range(num_qubits))
        else:
            self.qubit_indices = qubit_indices
            
        self.circuits = []
        self.transition_matrix = None
        self.basis_states = self._generate_basis_states()
        
    def _generate_basis_states(self) -> List[str]:
        """
        Generate all computational basis states for the selected qubits.
        
        Returns:
            List of binary strings representing all possible basis states.
        """
        n = len(self.qubit_indices)
        return [''.join(map(str, bits)) for bits in itertools.product([0, 1], repeat=n)]
    
    def add_circuit(self, circuit: QuantumCircuit):
        """
        Add a custom circuit to the benchmark set.
        
        Args:
            circuit: A quantum circuit to add to the benchmark
        """
        self.circuits.append(circuit)
        
    def generate_basis_circuits(self, add_cnot: bool = False, mapping: Dict[int, List[int]] = None) -> List[QuantumCircuit]:
        """
        Generate circuits to prepare each computational basis state.
        
        Args:
            add_cnot: Whether to add a CNOT gate to the circuit. default is False.
            mapping: A dictionary mapping the state of the physical qubits to the ancillary qubits.

        Returns:
            List of quantum circuits, each preparing a different basis state.
        """
        basis_circuits = []
        
        # Create a circuit for each basis state
        for basis_state in self.basis_states:
            qc = QuantumCircuit(self.num_qubits, len(self.qubit_indices))
            
            # Set qubits to the target basis state
            # Reverse the bit string to match conventional qubit ordering (MSB to LSB), this
            # is because the original qubit indices are in reverse order in qiskit, e.g. 01 means
            # q[1] = 0, q[0] = 1, but in the basis state, it is q[0] = 0, q[1] = 1, so we need to
            # reverse the bit string to get the correct basis state, e.g. 01 -> 10

            reversed_basis_state = basis_state[::-1]
            for idx, bit in zip(self.qubit_indices, reversed_basis_state):
                if bit == '1':
                    qc.x(idx)  # Flip qubit if basis state bit is 1
            
            if mapping:
                for target_qubit, ancillary_qubits in mapping.items():
                    physical_target_idx = self.qubit_indices.index(target_qubit)
                    for ancillary_qubit in ancillary_qubits:
                        physical_ancillary_idx = self.qubit_indices.index(ancillary_qubit)
                        qc.cx(physical_target_idx, physical_ancillary_idx)

            # Add measurements for the specified qubits
            for i, qubit_idx in enumerate(self.qubit_indices):
                qc.measure(qubit_idx, i)
                
            # Add this circuit to our collection
            basis_circuits.append(qc)
            self.circuits.append(qc)
        return basis_circuits
    
    def generate_logical_basis_circuits(self, mapping: Dict[int, List[int]], add_x_gates: bool = False) -> List[QuantumCircuit]:
        """
        Generate logical basis circuits based on the provided mapping.
        
        Args:
            mapping: A dictionary mapping the state of the physical qubits to the ancillary qubits.
            add_x_gates: Whether to add X gates to the ancillary qubits.
        """
        logical_basis_circuits = []
        
        self.logical_basis_states = [''.join(map(str, bits)) for bits in itertools.product([0, 1], repeat=len(mapping))]
        target_qubit_indices = list(mapping.keys())

        for logical_state in self.logical_basis_states:
            reversed_logical_state = logical_state[::-1]
            qc = QuantumCircuit(self.num_qubits, len(self.qubit_indices))
            for target_qubit, ancillary_qubits in mapping.items():
                # Find the index of the target qubit in the qubit_indices list
                target_idx = target_qubit_indices.index(target_qubit)
                physical_target_idx = self.qubit_indices.index(target_qubit)

                if reversed_logical_state[target_idx] == '1':
                    qc.x(physical_target_idx)
                # map the state of the target qubit to the ancillary qubits
                for ancillary_qubit in ancillary_qubits:
                    physical_ancillary_idx = self.qubit_indices.index(ancillary_qubit)
                    qc.cx(physical_target_idx, physical_ancillary_idx)

                    if add_x_gates:
                        qc.x(physical_ancillary_idx)
                    
                
            qc.measure_all(add_bits=False)
            logical_basis_circuits.append(qc)
            self.circuits.append(qc)

        return logical_basis_circuits

    

    def generate_transition_matrix(self, results: List[Dict[str, int]], mapping_state: Dict[str, str]=None, cnot_mapping: Dict[int, List[int]] = None, normalize: bool=True) -> np.ndarray:
        """
        Generate the readout error transition matrix from circuit execution results.
        
        Args:
            results: List of count dictionaries from executing the basis circuits
            mapping_state: A dictionary mapping the state of the physical states to the logical states.
            cnot_mapping: A dictionary mapping the state of the physical qubits to the ancillary qubits.
            normalize: Whether to normalize the transition matrix

        Returns:
            Transition matrix as a numpy array where T[i,j] represents the probability
            of measuring state j when preparing state i
        """
        n_states = len(self.basis_states)
        if mapping_state:
            n_states = len(mapping)

        transition_matrix = np.zeros((n_states, n_states))

        # For each prepared state (row in the transition matrix)
        if mapping_state:
            for i in range(n_states):
                counts = results[i]
                for measured_state, count in counts.items():
                    if measured_state in mapping:
                        transition_matrix[int(mapping[measured_state[::-1]], 2), i] = count    
        else:
            for idx, prepared_state in enumerate(self.basis_states):
                i = idx
                # Get the counts for this circuit execution
                counts = results[idx]
                total_shots = sum(counts.values())
                # Process the counts based on CNOT mapping if provided
                if cnot_mapping:
                   tmp_bitstring = list(prepared_state[::-1])  # Convert to list to make it mutable
                   for target_qubit, ancillary_qubits in cnot_mapping.items():
                       physical_target_idx = self.qubit_indices.index(target_qubit)
                       for ancillary_qubit in ancillary_qubits:
                           physical_ancillary_idx = self.qubit_indices.index(ancillary_qubit)
                           if tmp_bitstring[physical_target_idx] == '1':
                               # Flip the target bit if ancillary is 1
                               tmp_bitstring[physical_ancillary_idx] = '1' if tmp_bitstring[physical_ancillary_idx] == '0' else '0'
                   tmp_bitstring = ''.join(tmp_bitstring)
                   i = int(tmp_bitstring[::-1], 2)

                # For each possible measured state (column in the transition matrix)
                for j, measured_state in enumerate(self.basis_states):
                    # Get count for this measured state
                    count = counts.get(measured_state, j)
                    
                    # Calculate probability: count / total_shots
                    transition_matrix[j, i] = count
        
        if normalize:
            transition_matrix = transition_matrix / np.sum(transition_matrix, axis=0)
        
        self.transition_matrix = transition_matrix
        return transition_matrix
    
    

if __name__ == "__main__":
    # Create a benchmark for 2 qubits
    from qiskit_aer import AerSimulator
    from qiskit.transpiler.preset_passmanagers import generate_preset_pass_manager
    from qiskit_ibm_runtime.fake_provider import FakeGuadalupeV2

    # Initialize the fake backend
    backend = FakeGuadalupeV2()

    # generate a pass manager for compiling the benchmark circuits on the fake backend
    passmanager = generate_preset_pass_manager(optimization_level=1, backend=backend)

    # construct a benchmark circuit instance for 2 qubits
    benchmark = BenchmarkCircuit(num_qubits=2)

    # Generate all basis circuits
    benchmark_circuits = benchmark.generate_basis_circuits()
    shots = 1000

    # Transpile all benchmark circuits
    transpiled_circuits = passmanager.run(benchmark_circuits)

    # Initialize a list to store results for each circuit
    circuit_results = []

    # Run all transpiled circuits on the noisy simulator
    job = backend.run(transpiled_circuits, shots=shots)
    result = job.result()

    # Get the counts dictionary for each circuit
    for i in range(len(benchmark_circuits)):
        counts_dict = result.get_counts(i)
        circuit_results.append(counts_dict)
    
    # Generate the transition matrix
    transition_matrix = benchmark.generate_transition_matrix(circuit_results)
    print("Transition matrix:")
    print(transition_matrix)


    # ***************** test logical transition matrix *****************
    mapping = {0: [1]}
    logical_basis_circuits = benchmark.generate_logical_basis_circuits(mapping, add_x_gates=False)
    transpiled_logical_circuits = passmanager.run(logical_basis_circuits)
    job = backend.run(transpiled_logical_circuits, shots=shots)
    result = job.result()
    logical_circuit_results = []

    # Get the counts dictionary for each circuit
    for i in range(len(logical_basis_circuits)):
        counts_dict = result.get_counts(i)
        logical_circuit_results.append(counts_dict)
        print(counts_dict)
        
    # Generate the logical transition matrix
    logical_transition_matrix = benchmark.generate_transition_matrix(logical_circuit_results, mapping={'00': '0', '11': '1'})
    print("Logical transition matrix:")
    print(logical_transition_matrix)

    # ***************** test logical transition matrix with X gates *****************
    mapping = {0: [1]}
    logical_basis_circuits = benchmark.generate_logical_basis_circuits(mapping, add_x_gates=True)
    transpiled_logical_circuits = passmanager.run(logical_basis_circuits)
    job = backend.run(transpiled_logical_circuits, shots=shots)
    result = job.result()
    logical_circuit_results = []

    # Get the counts dictionary for each circuit
    for i in range(len(logical_basis_circuits)):
        counts_dict = result.get_counts(i)
        logical_circuit_results.append(counts_dict)
        print(counts_dict)
        
    # Generate the logical transition matrix
    logical_transition_matrix = benchmark.generate_transition_matrix(logical_circuit_results, mapping={'01': '0', '10': '1'})
    print("Logical transition matrix:")
    print(logical_transition_matrix)
